"""
AI对战模式 - 支持AI代理与玩家对战
"""
import pygame
import sys
from game import Game
from game_ai import GameAI, RandomAgent, TrainedAgent
from settings import *

# 使用向量进行精确的数学计算
vec = pygame.math.Vector2


class AIVsPlayerMode:
    """AI对战模式类"""
    
    def __init__(self, main_game):
        """初始化AI对战模式"""
        self.main_game = main_game
        self.game = main_game.game
        
        # AI相关
        self.ai_enabled = False
        self.ai_agent = None
        self.ai_think_time = 0  # AI思考时间（毫秒）
        self.ai_action_interval = 200  # AI动作间隔（毫秒）
        self.last_ai_action_time = 0
        
        # UI元素
        self.font = pygame.font.Font(None, 24)
        self.ai_status_text = "AI: 禁用"
        
    def toggle_ai(self):
        """切换AI状态"""
        self.ai_enabled = not self.ai_enabled
        
        if self.ai_enabled:
            # 启用AI
            self.ai_agent = GameAI(render_mode=None)  # 不渲染AI环境
            self.ai_status_text = "AI: 启用 (随机)"
            print("AI已启用")
        else:
            # 禁用AI
            self.ai_agent = None
            self.ai_status_text = "AI: 禁用"
            print("AI已禁用")
    
    def set_ai_model(self, model_path=None):
        """设置AI模型"""
        if model_path:
            # 使用训练好的模型
            self.ai_agent = GameAI(render_mode=None)
            self.ai_agent.agent = TrainedAgent(model_path)
            self.ai_status_text = f"AI: 训练模型 ({model_path})"
            print(f"已加载AI模型: {model_path}")
        else:
            # 使用随机AI
            self.ai_agent = GameAI(render_mode=None)
            self.ai_status_text = "AI: 随机"
            print("已设置为随机AI")
        
        self.ai_enabled = True
    
    def update(self, dt):
        """更新AI状态"""
        if not self.ai_enabled or not self.ai_agent:
            return
        
        # 检查是否到了AI执行动作的时间
        current_time = pygame.time.get_ticks()
        if current_time - self.last_ai_action_time < self.ai_action_interval:
            return
        
        # 检查游戏是否正在进行中
        if not hasattr(self.game, 'player') or not self.game.player or not self.game.player.alive():
            return
        
        # 检查是否有野怪
        if len(self.game.monsters) == 0:
            return
        
        # 执行AI动作
        self.execute_ai_action()
        self.last_ai_action_time = current_time
    
    def execute_ai_action(self):
        """执行AI动作"""
        # 获取当前游戏状态
        observation = self.get_game_observation()
        
        # AI选择动作
        action = self.ai_agent.agent.select_action(observation)
        
        # 执行动作
        self.execute_action(action)
    
    def get_game_observation(self):
        """获取当前游戏状态的观察"""
        # 创建一个简化的观察，与AI环境中的观察格式一致
        obs = [0] * 133
        
        # 玩家位置 (0-1)
        if hasattr(self.game, 'player') and self.game.player:
            obs[0] = self.game.player.pos.x
            obs[1] = self.game.player.pos.y
            obs[2] = self.game.player.health
        
        # 野怪信息 (3-32)
        monster_start = 3
        # 获取野怪列表
        monsters_list = list(self.game.monsters.sprites())
        for i, monster in enumerate(monsters_list[:10]):  # 最多10个野怪
            idx = monster_start + i * 3
            obs[idx] = monster.pos.x
            obs[idx + 1] = monster.pos.y
            obs[idx + 2] = monster.health
        
        # 子弹信息 (33-92)
        bullet_start = 33
        # 获取子弹列表
        bullets_list = list(self.game.bullets.sprites())
        for i, bullet in enumerate(bullets_list[:20]):  # 最多20个子弹
            idx = bullet_start + i * 3
            obs[idx] = bullet.pos.x
            obs[idx + 1] = bullet.pos.y
            # 将方向转换为角度
            angle = pygame.math.Vector2(bullet.vel).angle_to((1, 0))
            obs[idx + 2] = angle
        
        # 障碍物信息 (93-132)
        obstacle_start = 93
        # 获取障碍物列表
        obstacles_list = list(self.game.obstacles.sprites())
        for i, obstacle in enumerate(obstacles_list[:20]):  # 最多20个障碍物
            idx = obstacle_start + i * 2
            obs[idx] = obstacle.rect.x
            obs[idx + 1] = obstacle.rect.y
        
        return obs
    
    def execute_action(self, action):
        """执行AI选择的动作"""
        if not hasattr(self.game, 'player') or not self.game.player:
            return
        
        player = self.game.player
        
        if action < 9:  # 移动动作
            # 模拟键盘按键
            if action == 0:  # 不移动
                player.vel = vec(0, 0)
            elif action == 1:  # 向上
                player.vel = vec(0, -1)
            elif action == 2:  # 向下
                player.vel = vec(0, 1)
            elif action == 3:  # 向左
                player.vel = vec(-1, 0)
            elif action == 4:  # 向右
                player.vel = vec(1, 0)
            elif action == 5:  # 向左上
                player.vel = vec(-1, -1).normalize()
            elif action == 6:  # 向右上
                player.vel = vec(1, -1).normalize()
            elif action == 7:  # 向左下
                player.vel = vec(-1, 1).normalize()
            elif action == 8:  # 向右下
                player.vel = vec(1, 1).normalize()
            
            # 应用速度
            if action != 0:
                player.vel *= player.speed
        
        elif 9 <= action < 20:  # 射击动作
            # 计算射击方向
            shoot_angle = (action - 9) * (360 / 10)  # 每36度一个方向
            shoot_dir = vec(1, 0).rotate(shoot_angle)
            
            # 直接使用方向向量，不需要计算目标位置
            # 创建一个远处的目标位置，但确保方向正确
            target_pos = player.pos + shoot_dir * 1000  # 使用更远的距离确保方向准确
            
            # 执行射击
            player.shoot(target_pos)
    
    def draw_ui(self):
        """绘制AI状态UI"""
        # 绘制AI状态文本
        text_surface = self.font.render(self.ai_status_text, True, WHITE)
        self.game.screen.blit(text_surface, (10, 10))
        
        # 绘制控制提示
        help_text = "按A键切换AI状态"
        help_surface = self.font.render(help_text, True, WHITE)
        self.game.screen.blit(help_surface, (10, 40))
    
    def handle_event(self, event):
        """处理事件"""
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_a:
                self.toggle_ai()


class AIVsPlayerMain:
    """AI对战模式主程序"""
    
    def __init__(self):
        """初始化AI对战模式"""
        pygame.init()
        self.game = Game()
        self.ai_mode = AIVsPlayerMode(self)
        
        # 设置游戏状态
        self.game.state = 'game'
        self.game.new()
        
        # 游戏循环控制
        self.running = True
        self.clock = pygame.time.Clock()
    
    def run(self):
        """运行AI对战模式"""
        while self.running:
            dt = self.clock.tick(FPS) / 1000.0
            
            # 处理事件
            self.handle_events()
            
            # 更新游戏
            self.game.update()
            
            # 更新AI
            self.ai_mode.update(dt)
            
            # 绘制
            self.draw()
        
        pygame.quit()
        sys.exit()
    
    def handle_events(self):
        """处理事件"""
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False
            
            # 处理AI模式事件
            self.ai_mode.handle_event(event)
            
            # 处理游戏事件
            self.game.game_events(event)
    
    def draw(self):
        """绘制游戏"""
        self.game.draw()
        
        # 绘制AI状态UI
        self.ai_mode.draw_ui()
        
        pygame.display.flip()


if __name__ == '__main__':
    """程序入口点"""
    ai_vs_player = AIVsPlayerMain()
    ai_vs_player.run()